import os
import base64
import json
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding, hashes
from cryptography.hazmat.backends import default_backend


class AES256Cipher:


    def __init__(self, key=None):
        """
        初始化加密工具

        :param key: 可选的密钥(32字节)，如果为None则自动生成
        """
        if key is None:
            # 生成随机密钥
            self.key = os.urandom(32)
        else:
            if len(key) != 32:
                raise ValueError("AES256 需要 32 字节的密钥")
            self.key = key

    @staticmethod
    def generate_key():
        """生成一个随机的 AES-256 密钥"""
        return os.urandom(32)

    @staticmethod
    def derive_key(password: str):
        """
        从密码直接派生密钥 (简单SHA256哈希)

        :param password: 密码字符串
        :return: 派生密钥
        """
        digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
        digest.update(password.encode('utf-8'))
        return digest.finalize()

    def encrypt_cbc(self, plaintext: str):
        """
        AES-256 CBC 模式加密

        :param plaintext: 要加密的明文(字符串)
        :return: 加密后的数据字符串(包含密文和IV)
        """
        # 生成随机 IV
        iv = os.urandom(16)

        # 创建加密器
        padder = padding.PKCS7(128).padder()
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.CBC(iv),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()

        # 填充并加密数据
        padded_data = padder.update(plaintext.encode('utf-8')) + padder.finalize()
        ciphertext = encryptor.update(padded_data) + encryptor.finalize()

        # 组合结果为一个字符串 (格式: iv.ciphertext)
        return f"{base64.b64encode(iv).decode('utf-8')}.{base64.b64encode(ciphertext).decode('utf-8')}"

    def decrypt_cbc(self, encrypted_str: str):
        """
        AES-256 CBC 模式解密

        :param encrypted_str: 加密的数据字符串(格式: iv.ciphertext)
        :return: 解密后的明文
        """
        # 分割字符串
        parts = encrypted_str.split('.')
        if len(parts) != 2:
            raise ValueError("无效的加密字符串格式")

        iv_b64, ciphertext_b64 = parts

        # 解码 base64
        ciphertext = base64.b64decode(ciphertext_b64)
        iv = base64.b64decode(iv_b64)

        # 创建解密器
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.CBC(iv),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()
        unpadder = padding.PKCS7(128).unpadder()

        # 解密并去除填充
        padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()
        plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()

        return plaintext.decode('utf-8')

    def encrypt_gcm(self, plaintext: str, associated_data: str = None):
        """
        AES-256 GCM 模式加密(提供认证)

        :param plaintext: 要加密的明文(字符串)
        :return: 加密后的数据字符串(包含密文、nonce和tag)
        :param associated_data: 可选的关联数据(用于认证但不加密)
        """
        # 生成随机 nonce(对于 GCM 推荐 12 字节)
        nonce = os.urandom(12)

        # 创建加密器
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.GCM(nonce),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()

        # 添加关联数据(如果有)
        if associated_data:
            encryptor.authenticate_additional_data(associated_data.encode('utf-8'))

        # 加密数据
        ciphertext = encryptor.update(plaintext.encode('utf-8')) + encryptor.finalize()

        # 组合结果为一个字符串 (格式: nonce.ciphertext.tag)
        return f"{base64.b64encode(nonce).decode('utf-8')}.{base64.b64encode(ciphertext).decode('utf-8')}.{base64.b64encode(encryptor.tag).decode('utf-8')}"

    def decrypt_gcm(self, encrypted_str: str, associated_data: str = None):
        """
        AES-256 GCM 模式解密

        :param encrypted_str: 加密的数据字符串(格式: nonce.ciphertext.tag)
        :param associated_data: 可选的关联数据(必须与加密时相同)
        :return: 解密后的明文
        :raises ValueError: 如果认证失败
        """
        # 分割字符串
        parts = encrypted_str.split('.')
        if len(parts) != 3:
            raise ValueError("无效的加密字符串格式")

        nonce_b64, ciphertext_b64, tag_b64 = parts

        # 解码 base64
        ciphertext = base64.b64decode(ciphertext_b64)
        nonce = base64.b64decode(nonce_b64)
        tag = base64.b64decode(tag_b64)

        # 创建解密器
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.GCM(nonce, tag),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()

        # 添加关联数据(如果有)
        if associated_data:
            decryptor.authenticate_additional_data(associated_data.encode('utf-8'))

        # 解密并验证
        try:
            plaintext = decryptor.update(ciphertext) + decryptor.finalize()
            return plaintext.decode('utf-8')
        except Exception as e:
            raise ValueError("认证失败 - 数据可能被篡改") from e

    # 文件加密/解密方法保持不变...

# 使用示例
# 2. 使用密码派生密钥
# password = "szsk123"
# key = AES256Cipher.derive_key(password)
# cipher = AES256Cipher(key)
#
# # CBC 模式
# encrypted_str = cipher.encrypt_cbc("这是一个秘密消息")
# print("加密结果:", encrypted_str)
# plaintext = cipher.decrypt_cbc("Ku/aa6/v9m8TNsu+BZnHig==.vZB6s0geTf8Eehs7ju+ClA==")
# print("解密结果:", plaintext)

# # GCM 模式
# encrypted_str = cipher.encrypt_gcm("这是一个需要认证的秘密消息", "关联数据")
# print("\n加密结果:", encrypted_str)
# plaintext = cipher.decrypt_gcm(encrypted_str, "关联数据")
# print("解密结果:", plaintext)